import torch
from utils.util import LD_param as get_res
from utils.util import AverageMeter,save_csv,mergetotxt,roc_plot
import os
import numpy as np

class TestApi():
    def __init__(self,opt):
        self.results_dir = opt.results_dir
        self.name = opt.name
        self.version = opt.version
        self.phase = opt.phase
        self.epoch = opt.epoch
        self.arch_name = opt.arch_name
        self.arch_type = opt.arch_type
        self.input_modality = opt.input_modality
        pass
    def init(self,controller,loss_names):
        self.TP,self.TN,self.FP,self.FN = {},{},{},{}
        self.TP['all'],self.TN['all'],self.FP['all'],self.FN['all']  = 0,0,0,0
        self.cat_name = None
        self.loss_names = loss_names
        if controller == 'test':
            self.TP['物理'],self.TN['物理'],self.FP['物理'],self.FN['物理'], = 0,0,0,0
            self.groud_truth = []
            self.prediction = []
            self.features = []
            self.cats = []

            
    def record(self,records):#label,out,controller,cat,cat_name):
        if "cat_name" in records.keys():
            self.cat_name = records["cat_name"]
        ans = torch.max(records["out"]['prediction'],axis=1).indices.cpu()
        self.TP['all'] += torch.where((ans==1) &(records["label"] == 1))[0].shape[0]
        self.TN['all'] += torch.where((ans==0) &(records["label"] == 0))[0].shape[0]
        self.FP['all'] += torch.where((ans==1) &(records["label"] == 0))[0].shape[0]
        self.FN['all'] += torch.where((ans==0) &(records["label"] == 1))[0].shape[0]
        if  records["controller"] == 'test' and self.cat_name is not None and records["cat"] is not None:
            # self.prediction += ans.tolist()
            # self.prediction.append(records["out"]['prediction'].max().cpu().detach().numpy())
            for fea in records["out"]['feature']:
                self.features.append(fea.squeeze(0).contiguous().cpu().detach().numpy())
            self.cats += records["cat"].tolist()
            self.groud_truth += records["label"].tolist()
            self.prediction +=records["out"]['prediction'][:,1].contiguous().cpu().detach().numpy().tolist()
            for key in self.cat_name.keys():
                if not key in self.TP.keys():
                    self.TP[key]=0
                    self.TN[key]=0
                    self.FP[key]=0
                    self.FN[key]=0
                self.TP[key] += torch.where((records["cat"] == self.cat_name[key])&(ans==1) &(records["label"] == 1))[0].shape[0]
                self.TN[key] += torch.where((records["cat"] == self.cat_name[key])&(ans==0) &(records["label"] == 0))[0].shape[0]
                self.FP[key] += torch.where((records["cat"] == self.cat_name[key])&(ans==1) &(records["label"] == 0))[0].shape[0]
                self.FN[key] += torch.where((records["cat"] == self.cat_name[key])&(ans==0) &(records["label"] == 1))[0].shape[0]
                # if records["cat"] >= 0 and records["cat"] <=4:
            self.TP['物理'] += torch.where((ans==1) &(records["label"] == 1)&(records["cat"] >= 0)&(records["cat"] <=4))[0].shape[0]
            self.TN['物理'] += torch.where((ans==0) &(records["label"] == 0)&(records["cat"] >= 0)&(records["cat"] <=4))[0].shape[0]
            self.FP['物理'] += torch.where((ans==1) &(records["label"] == 0)&(records["cat"] >= 0)&(records["cat"] <=4))[0].shape[0]
            self.FN['物理'] += torch.where((ans==0) &(records["label"] == 1)&(records["cat"] >= 0)&(records["cat"] <=4))[0].shape[0]
        return get_res(self.TP,self.TN,self.FP,self.FN,2)
    
    def statistic(self,controller):
        param = get_res(self.TP,self.TN,self.FP,self.FN,2)
        if controller == 'test':
            
            self.cats = np.array(self.cats)
            self.cats[np.where(np.array(self.cats)<=4)] = 0
            self.cats[np.where(np.array(self.cats)==5)] = 1
            self.cats[np.where(np.array(self.cats)==6)] = 2
            
            
            self.features = np.array(self.features)
            # x,y = self.features[:,-1],self.features[:,-2]
            arg1=50
            arg2=1000
            fea1_mean = np.mean(self.features[np.where(np.array(self.cats)==0),:][0],0)
            fea2_mean = np.mean(self.features[np.where(np.array(self.cats)==1),:][0],0)
            fea3_mean = np.mean(self.features[np.where(np.array(self.cats)==2),:][0],0)
            fea_mean = np.abs(fea1_mean-fea2_mean)+np.abs(fea1_mean-fea3_mean)+np.abs(fea2_mean-fea3_mean)
            fea_mean -= np.var(self.features[np.where(np.array(self.cats)==0),:][0],0)*arg1
            fea_mean -= np.var(self.features[np.where(np.array(self.cats)==1),:][0],0)*arg1
            fea_mean -= np.var(self.features[np.where(np.array(self.cats)==2),:][0],0)*arg1
            fea_mean += arg2
            pos = -1
            while (fea_mean[pos] >990) or (np.sum(self.features[:,pos])<1000):
                pos -=1
            indx = pos
            pos-=1
            while (fea_mean[pos] >990) or (np.sum(self.features[:,pos])<1000):
                pos -=1
            indy = pos
            # indx,indy =np.argsort(fea_mean)[posx],np.argsort(fea_mean)[posy]
            x,y =  self.features[:,indx],self.features[:,indy]

            import matplotlib.pyplot as plt
            
            plt.cla()
            plt.scatter(x[np.where(self.cats==0)],y[np.where(self.cats==0)],0.5,c='g',label='Features of Physical Domain')
            plt.scatter(x[np.where(self.cats==1)],y[np.where(self.cats==1)],0.5,c='r',label='Features of Geniune Domain')
            plt.scatter(x[np.where(self.cats==2)],y[np.where(self.cats==2)],0.5,c='b',label='Features of Digital Domain')
            plt.legend()
            plt.title('Feature disturbution on 2 dimension')
            plt.xlabel('The '+str(indx)+'th dimension feature')
            plt.ylabel('The '+str(indx)+'th dimension feature')
            
            csvdir = os.path.join(self.results_dir,self.name+'_v'+self.version,self.phase+'_'+self.epoch)
            plt.savefig(os.path.join(csvdir,"FeatureSpace.png"))
            
            csvfile = '{}_{}_{}_{}'.format(self.arch_name,self.arch_type,self.input_modality,'v'+self.version)
            sheet_name = '{}{}'.format(self.arch_name+self.arch_type,self.input_modality.upper())
            if not '物理' in self.TP.keys():
                save_csv(param,csvdir,csvfile,sheet_name,self.loss_names,list(self.cat_name.keys())+['all','物理'])
            elif self.cat_name is not None and '物理' in self.TP.keys():
                save_csv(param,csvdir,csvfile,sheet_name,self.loss_names,list(self.cat_name.keys())+['all','物理'])
            elif self.cat_name is not None and not '物理' in self.TP.keys():
                save_csv(param,csvdir,csvfile,sheet_name,self.loss_names,list(self.cat_name.keys())+['all'])
            else:
                save_csv(param,csvdir,csvfile,sheet_name,self.loss_names,['all'])

        # roc_plot(self.groud_truth,self.prediction,'test.png','/home/yalin.huang/test/results/'+self.name)
            mergetotxt(self.groud_truth,self.prediction,csvdir,self.arch_name,self.version) 
            self.TP['物理'],self.TN['物理'],self.FP['物理'],self.FN['物理'], = 0,0,0,0
        self.TP,self.TN,self.FP,self.FN = {},{},{},{}
        self.TP['all'],self.TN['all'],self.FP['all'],self.FN['all']  = 0,0,0,0